import pandas as pd
import sys
from PySide6.QtGui import QKeySequence, QShortcut
from PySide6.QtWidgets import QApplication, QWidget, QVBoxLayout, QLabel, QPushButton, QLineEdit, QPlainTextEdit

'''
按Q是白字快捷键
按W是黑子快捷键
正常点叉叉退出会保存，写入文件的
自动从未标记的地方开始
'''


class DataAnnotator:
    def __init__(self, dataset, current_index=0):
        self.dataset = dataset
        self.current_index = current_index
        self.label_value = None

        self.app = QApplication(sys.argv)
        self.window = QWidget()
        self.window.setGeometry(100, 100, 800, 600)
        self.layout = QVBoxLayout()

        # 创建一个 QLabel 用于显示提示信息
        self.label_prompt = QLabel("快捷键Q白子-1  |  快捷键W黑子-0  |  点叉退出自动保存")
        self.label_prompt.setGeometry(100, 20, 200, 30)  # 设置位置和大小的示例
        self.layout.addWidget(self.label_prompt)

        # 创建一个 QLineEdit 用于显示数据集中的文本
        self.text_display = QPlainTextEdit(self.dataset['txt'].iloc[self.current_index])
        self.text_display.setReadOnly(True)
        self.text_display.setMinimumWidth(400)
        self.text_display.setMinimumHeight(400)
        self.text_display.setStyleSheet("font-size: 30px;")
        self.layout.addWidget(self.text_display)

        # 创建两个 QPushButton 用于标注标签
        self.button_1 = QPushButton('白子')
        self.button_1.clicked.connect(self.set_label_1)
        self.layout.addWidget(self.button_1)

        self.button_0 = QPushButton('黑子')
        self.button_0.clicked.connect(self.set_label_0)
        self.layout.addWidget(self.button_0)

        self.window.setLayout(self.layout)
        self.window.setWindowTitle('NLP Data Label (作者：凯哥)')

        # 连接 aboutToQuit 信号到 save_to_file 方法
        self.app.aboutToQuit.connect(self.save_to_file)

        # 创建一个 QShortcut 对象,将快捷键与按钮的 clicked 信号关联
        self.my_shortcut = QShortcut(QKeySequence("Q"), self.window).activated.connect(self.button_1.click)
        self.my_shortcut1 = QShortcut(QKeySequence("W"), self.window).activated.connect(self.button_0.click)

        self.window.show()

        self.app.exec()

    def set_label_1(self):
        self.label_value = 1
        self.next_row()

    def set_label_0(self):
        self.label_value = 0
        self.next_row()

    def next_row(self):
        # 保存当前数据
        self.dataset.loc[self.current_index, 'label'] = self.label_value
        # 加载下一条数据
        self.current_index += 1
        if self.current_index < len(self.dataset):
            self.text_display.setPlainText(self.dataset['txt'].iloc[self.current_index])
        else:
            self.app.quit()

    def save_to_file(self):
        # 保存数据集到文件
        self.dataset.to_csv('../csv_dir/annotated_data.csv', index=False)
        print("数据已保存到 annotated_data.csv")


def load_last_date():
    # 加载数据集
    file_path = '../csv_dir/annotated_data.csv'
    dataset = pd.read_csv(file_path)
    # 得到未标注数据的行号
    rows_with_na = dataset[dataset.isna().any(axis=1)]

    # print(dataset.shape)

    # 创建数据标注工具
    annotator = DataAnnotator(dataset, int(*rows_with_na[:1]['ID'].values))


load_last_date()
